"""
ES和扩展F信号检测模型训练脚本
基于Vision Transformer，针对MX450显卡优化
"""

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR, ReduceLROnPlateau
import numpy as np
import matplotlib.pyplot as plt
import os
import time

from datetime import datetime
import json

from dataset import create_data_loaders
from vit_model import create_model
from label_mapping import ID_TO_DESCRIPTION

class Trainer:
    def __init__(self, model, train_loader, val_loader, device, config):
        self.model = model.to(device)
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.device = device
        self.config = config
        
        # 损失函数和优化器
        self.criterion = nn.CrossEntropyLoss()
        self.optimizer = optim.AdamW(
            model.parameters(), 
            lr=config['learning_rate'],
            weight_decay=config['weight_decay']
        )
        
        # 学习率调度器
        if config['scheduler'] == 'cosine':
            self.scheduler = CosineAnnealingLR(
                self.optimizer, 
                T_max=config['epochs'],
                eta_min=config['learning_rate'] * 0.01
            )
        else:
            self.scheduler = ReduceLROnPlateau(
                self.optimizer, 
                mode='min', 
                factor=0.5, 
                patience=5,
                verbose=True
            )
        
        # 训练历史
        self.train_losses = []
        self.train_accs = []
        self.val_losses = []
        self.val_accs = []
        self.best_val_acc = 0.0
        self.best_model_path = None
        
    def train_epoch(self):
        """训练一个epoch"""
        self.model.train()
        total_loss = 0.0
        correct = 0
        total = 0
        
        for batch_idx, (images, labels, _) in enumerate(self.train_loader):
            images, labels = images.to(self.device), labels.to(self.device)
            
            # 前向传播
            self.optimizer.zero_grad()
            outputs = self.model(images)
            loss = self.criterion(outputs, labels)
            
            # 反向传播
            loss.backward()
            
            # 梯度裁剪
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
            
            self.optimizer.step()
            
            # 统计
            total_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            
            # 打印进度
            if batch_idx % 10 == 0:
                print(f'Batch {batch_idx}/{len(self.train_loader)}, '
                      f'Loss: {loss.item():.4f}, '
                      f'Acc: {100.*correct/total:.2f}%')
        
        avg_loss = total_loss / len(self.train_loader)
        accuracy = 100. * correct / total
        
        return avg_loss, accuracy
    
    def validate(self):
        """验证模型"""
        self.model.eval()
        total_loss = 0.0
        correct = 0
        total = 0
        class_correct = [0] * 5
        class_total = [0] * 5
        
        with torch.no_grad():
            for images, labels, _ in self.val_loader:
                images, labels = images.to(self.device), labels.to(self.device)
                
                outputs = self.model(images)
                loss = self.criterion(outputs, labels)
                
                total_loss += loss.item()
                _, predicted = torch.max(outputs, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
                
                # 每个类别的准确率
                for i in range(labels.size(0)):
                    label = labels[i]
                    class_correct[label] += (predicted[i] == label).item()
                    class_total[label] += 1
        
        avg_loss = total_loss / len(self.val_loader)
        accuracy = 100. * correct / total
        
        # 打印每个类别的准确率
        print("\n各类别准确率:")
        for i in range(5):
            if class_total[i] > 0:
                class_acc = 100. * class_correct[i] / class_total[i]
                print(f'{ID_TO_DESCRIPTION[i]}: {class_acc:.2f}% ({class_correct[i]}/{class_total[i]})')
        
        return avg_loss, accuracy
    
    def save_model(self, epoch, val_acc, is_best=False):
        """保存模型"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        if is_best:
            filename = f'best_model_{timestamp}.pth'
        else:
            filename = f'model_epoch_{epoch}_{timestamp}.pth'
        
        filepath = os.path.join('models', filename)
        os.makedirs('models', exist_ok=True)
        
        torch.save({
            'epoch': epoch,
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'val_acc': val_acc,
            'config': self.config,
            'model_info': self.model.get_model_info()
        }, filepath)
        
        if is_best:
            self.best_model_path = filepath
        
        print(f'模型已保存: {filepath}')
        return filepath
    
    def plot_training_history(self):
        """绘制训练历史"""
        plt.figure(figsize=(12, 4))
        
        # 损失曲线
        plt.subplot(1, 2, 1)
        plt.plot(self.train_losses, label='训练损失')
        plt.plot(self.val_losses, label='验证损失')
        plt.title('训练和验证损失')
        plt.xlabel('Epoch')
        plt.ylabel('Loss')
        plt.legend()
        plt.grid(True)
        
        # 准确率曲线
        plt.subplot(1, 2, 2)
        plt.plot(self.train_accs, label='训练准确率')
        plt.plot(self.val_accs, label='验证准确率')
        plt.title('训练和验证准确率')
        plt.xlabel('Epoch')
        plt.ylabel('Accuracy (%)')
        plt.legend()
        plt.grid(True)
        
        plt.tight_layout()
        plt.savefig('training_history.png', dpi=300, bbox_inches='tight')
        plt.show()
    
    def train(self):
        """完整训练流程"""
        print("开始训练...")
        print(f"设备: {self.device}")
        print(f"训练样本: {len(self.train_loader.dataset)}")
        print(f"验证样本: {len(self.val_loader.dataset)}")
        print("-" * 60)
        
        for epoch in range(self.config['epochs']):
            start_time = time.time()
            
            print(f'Epoch {epoch+1}/{self.config["epochs"]}')
            
            # 训练
            train_loss, train_acc = self.train_epoch()
            
            # 验证
            val_loss, val_acc = self.validate()
            
            # 更新学习率
            if self.config['scheduler'] == 'cosine':
                self.scheduler.step()
            else:
                self.scheduler.step(val_loss)
            
            # 记录历史
            self.train_losses.append(train_loss)
            self.train_accs.append(train_acc)
            self.val_losses.append(val_loss)
            self.val_accs.append(val_acc)
            
            # 保存最佳模型
            is_best = val_acc > self.best_val_acc
            if is_best:
                self.best_val_acc = val_acc
                self.save_model(epoch+1, val_acc, is_best=True)
            
            # 定期保存模型
            if (epoch + 1) % 10 == 0:
                self.save_model(epoch+1, val_acc, is_best=False)
            
            epoch_time = time.time() - start_time
            current_lr = self.optimizer.param_groups[0]['lr']
            
            print(f'训练损失: {train_loss:.4f}, 训练准确率: {train_acc:.2f}%')
            print(f'验证损失: {val_loss:.4f}, 验证准确率: {val_acc:.2f}%')
            print(f'学习率: {current_lr:.6f}, 用时: {epoch_time:.2f}s')
            print(f'最佳验证准确率: {self.best_val_acc:.2f}%')
            print("-" * 60)
        
        print("训练完成!")
        print(f"最佳验证准确率: {self.best_val_acc:.2f}%")
        print(f"最佳模型保存在: {self.best_model_path}")
        
        # 绘制训练历史
        self.plot_training_history()
        
        return self.best_model_path

def main():
    """主函数"""
    # 配置参数
    config = {
        'batch_size': 16,        # 适合MX450的批次大小
        'learning_rate': 1e-4,   # 较小的学习率
        'weight_decay': 1e-4,
        'epochs': 50,
        'scheduler': 'cosine',   # 'cosine' 或 'plateau'
        'image_size': 224,
        'train_split': 0.8
    }
    
    # 设备检查
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"使用设备: {device}")
    
    if torch.cuda.is_available():
        print(f"GPU: {torch.cuda.get_device_name(0)}")
        print(f"显存: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
    
    # 创建数据加载器
    print("加载数据...")
    train_loader, val_loader = create_data_loaders(
        image_dir=".",
        labels_file="0818.txt",
        image_names_file="mingzi.txt",
        batch_size=config['batch_size'],
        train_split=config['train_split'],
        image_size=config['image_size']
    )
    
    # 创建模型
    print("创建模型...")
    model = create_model(num_classes=5)
    
    # 打印模型信息
    model_info = model.get_model_info()
    print(f"模型参数量: {model_info['total_parameters']:,}")
    
    # 创建训练器并开始训练
    trainer = Trainer(model, train_loader, val_loader, device, config)
    best_model_path = trainer.train()
    
    # 保存配置
    with open('training_config.json', 'w', encoding='utf-8') as f:
        json.dump(config, f, indent=2, ensure_ascii=False)
    
    print(f"训练配置已保存到: training_config.json")
    print(f"最佳模型路径: {best_model_path}")

if __name__ == "__main__":
    main()